import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
# 加载 MNIST 数据集
trainset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform)
# 创建数据加载器
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

# 定义神经网络类


class Net(nn.Module):
    def __init__(self):
        """
        初始化神经网络
        """
        super(Net, self).__init__()
        # 创建一个二维卷积层，输入通道为1说明是灰度图，输出通道数为32， 卷积核大小为3x3，填充为1说明在边缘进行卷积是会填充
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        # 创建一个二维卷积层，输入通道为32， 输出通道为64， 卷积核大小为3x3，填充为1
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        # 创建一个最大池化层，池化核大小为2x2，步长为2
        self.pool = nn.MaxPool2d(2, 2)
        # 创建第一个全连接层，输入的维度为64*7*7，输出维度为128
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        # 创建第二个全连接层，输入维度为128，输出维度为10（10个数字也就是10个分类）
        self.fc2 = nn.Linear(128, 10)
        # 创建一个丢弃层，丢弃的概率为0.5，防止模型过拟合，提高模型的泛化能力
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        """
        前向计算，
        :param x: 输入到神经网络的数据，这里支灰度图片
        :return: 经过神经网络前向计算后得到的输出结果（每个分类结果的概率）
        """
        # 第一个卷积操作
        x = self.pool(torch.relu(self.conv1(x)))
        # 第二个卷积操作
        x = self.pool(torch.relu(self.conv2(x)))
        # 将池化操作后的数据进项形状重塑为我们指定的形状（第一个全连接层输入的形状）
        x = x.view(-1, 64 * 7 * 7)
        # 将重塑后的数据传入第一个全连接层进行处理，并且应用relu激活函数处理
        x = torch.relu(self.fc1(x))
        # 经过丢弃层处理
        x = self.dropout(x)
        # 再经过第二个全连接层进行处理
        x = self.fc2(x)
        return x


# 创建模型实例
model = Net()
# 定义损失函数，pytorch中常用于多分类任务的损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器，Adam优化器是基于梯度的优化算法，设置了可学习参数model.parameters()和学习率lr=0.001
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练循环，开始每一轮的训练，这里我们设置最大训练轮次max_epoch为10
for epoch in range(10):
    running_loss = 0.0  # 设置一个损失值
    # 这里可以看到最大迭代次数max_iteration就是数据的大小，说明批大小batch_size也就是整个数据大小
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data  # 获取数据的输入和标签
        optimizer.zero_grad()  # 查资料说pytorch的梯度会自动累加，这是个优化器的梯度清零操作
        # 输入数据inputs经过前向传播得到输出outputs
        outputs = model(inputs)
        # 计算模型的的预测输出outputs和真实标签labels之间的损失值
        loss = criterion(outputs, labels)
        # 前面没有定义反向传播函数，查资料说这是pytorch中的损失函数返回类型包括了反向传播，意味着调用backward
        # 会触发pytorch的自动求导机制，自动计算每个可学习参数的梯度
        loss.backward()
        # 优化器每个可学习参数更新
        optimizer.step()
        # 将每次训练后的损失累加起来，为后面检查点100求该检查点的平均损失
        running_loss += loss.item()
        if i % 100 == 99:
            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 100}')
            running_loss = 0.0
# 在测试集上评估模型
correct = 0  # 准确率
total = 0  #
with torch.no_grad():  # 暂时关闭梯度计算
    for data in testloader:  # 对于测试机中每个数据
        images, labels = data  # 获取灰度图片的对应标签
        outputs = model(images)  # 模型预测
        _, predicted = torch.max(outputs.data, 1) # 获取模型预测结果中概率最大的那个类别的下标
        total += labels.size(0)  # 统计标签的数量
        correct += (predicted == labels).sum().item()  # 统计正确预测的数量
# 打印模型在测试集的准确率
print(f'Accuracy on test set: {100 * correct / total}%')
